from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QPushButton,
                             QLabel, QProgressBar, QFileDialog, QMessageBox)
from PyQt5.QtCore import QThread, pyqtSignal


class FrameExtractorThread(QThread):
    progress_updated = pyqtSignal(int)
    finished = pyqtSignal()
    error_occurred = pyqtSignal(str)

    def __init__(self, video_path, output_dir):
        super().__init__()
        self.video_path = video_path
        self.output_dir = output_dir

    def run(self):
        try:
            import cv2
            import os

            # 确保输出目录存在
            os.makedirs(self.output_dir, exist_ok=True)

            # 打开视频文件
            cap = cv2.VideoCapture(self.video_path)
            if not cap.isOpened():
                raise Exception("无法打开视频文件")

            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            current_frame = 0

            while True:
                ret, frame = cap.read()
                if not ret:
                    break

                # 保存帧为图片
                frame_path = os.path.join(self.output_dir, f"frame_{current_frame:04d}.jpg")
                cv2.imwrite(frame_path, frame)

                current_frame += 1
                progress = int((current_frame / total_frames) * 100)
                self.progress_updated.emit(progress)

            cap.release()
            self.finished.emit()
        except Exception as e:
            self.error_occurred.emit(str(e))


class VideoExtractorUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("视频解帧工具")
        self.setGeometry(100, 100, 400, 200)

        # 初始化UI元素
        self.video_path_a = None
        self.init_ui()

    def init_ui(self):
        # 主窗口布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout()

        # 视频选择按钮
        self.video_btn = QPushButton("选择视频文件")
        self.video_btn.clicked.connect(self.select_video)
        layout.addWidget(self.video_btn)

        # 解帧按钮
        self.extract_btn = QPushButton("开始解帧")
        self.extract_btn.clicked.connect(self.extract_frames)
        layout.addWidget(self.extract_btn)

        # 进度条
        self.progress_bar = QProgressBar()
        layout.addWidget(self.progress_bar)

        central_widget.setLayout(layout)

    def select_video(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择视频文件", "", "视频文件 (*.mp4 *.avi)")
        if file_path:
            self.video_path_a = file_path
            QMessageBox.information(self, "提示", f"已选择视频: {file_path}")

    def extract_frames(self):
        """异步执行视频解帧"""
        if not self.video_path_a:
            QMessageBox.warning(self, "警告", "请先选择视频文件")
            return

        # 弹出目录选择对话框
        output_dir = QFileDialog.getExistingDirectory(self, '选择帧输出目录')
        if not output_dir:
            return

        # 禁用按钮避免重复点击
        self.extract_btn.setEnabled(False)
        self.progress_bar.setValue(0)

        # 创建并启动解帧线程
        self.extract_thread = FrameExtractorThread(self.video_path_a, output_dir)
        self.extract_thread.progress_updated.connect(self.update_extract_progress)
        self.extract_thread.finished.connect(self.extract_complete)
        self.extract_thread.error_occurred.connect(self.extract_error)
        self.extract_thread.start()

    def update_extract_progress(self, value):
        self.progress_bar.setValue(value)

    def extract_complete(self):
        self.extract_btn.setEnabled(True)
        QMessageBox.information(self, "提示", "解帧完成！")

    def extract_error(self, error_msg):
        self.extract_btn.setEnabled(True)
        QMessageBox.critical(self, "错误", f"解帧失败: {error_msg}")


if __name__ == '__main__':
    app = QApplication([])
    window = VideoExtractorUI()
    window.show()
    app.exec_()
